import os
from re import T
import pandas as pd

from datetime import datetime
from WindPy import w

today = datetime.today().strftime('%Y-%m-%d')

def get_batch_wind_data(id_to_name):
    if not w.isconnected():
        w.start() 

    keys = list(id_to_name.keys())
    s = ','.join(keys)

    re = w.edb(
        s, 
        "1999-12-28", 
        # "2025-06-03"
        today
        )

    df = pd.DataFrame(
        re.Data,
        index=re.Codes,
        columns=re.Times,
    ).T
    df.columns = [id_to_name[k] for k in keys]
    return df

def get_data(freq = 'M', use_cache=True, sub_dir=False):
    save_cache = f"./data_cache/{freq}_raw_data_{today}.xlsx"
    if sub_dir:
        save_cache = f"./{save_cache}"
    
    if use_cache and os.path.exists(save_cache):
        print('Loading cache', save_cache)
        return pd.read_excel(save_cache, index_col=0)
    path = r"tjd_meta.xlsx"
    meta_df = pd.read_excel(path, index_col=0)

    freq_map = {'M': '月', 'Y':'年', 'Q':'季'}
    df_m = meta_df[meta_df['频率'] == freq_map[freq]]
    id_name = dict(zip(df_m['指标ID'].to_list(), df_m['指标名称'].to_list()))
    name_tjd_name = dict(zip(df_m['指标名称'].to_list(), df_m['腾景名称'].to_list()))

    data = get_batch_wind_data(id_name)
    data.columns = [name_tjd_name[k] for k in data.columns]

    data.to_excel(save_cache)
    return data


if __name__ == '__main__':
    m_raw_data = get_data(freq = 'M', use_cache=True)
    # get_data('Y')